
#---------------------------------------
#
#Script: mailbox.py
#Description: takes in a mailbox format text files and puts
#             relevant data into a SQL database (postgresql)
#Author: Alex Gourley (acgourley@ucdavis.edu)
#
#---------------------------------------

import sys
import email
import email.Errors
import _pg
import re
import time
import timezone
import popen2
import string 
import os



def getSqlFileText(sqlFile):
	"""given the name of a file that holds an sql query, find it,
	read it, and return the text.  It first looks for $SQLDIR from
	the environment and checks there for the file, then checks in ../sql/"""
	
	if os.environ.has_key("SQLDIR"):
		sqlpath = os.path.join(os.environ["SQLDIR"], sqlFile)
		if os.path.exists(path):
			return open(path).read()
	return open("../sql/" + sqlFile).read()
		

timeZoneTable = timezone.TimeZoneTable()

class DataError(Exception):
	def __init__(self, value):
		self.value = value

def getBody(message):
	if message.is_multipart():
		body = ""
		for submessage in message.get_payload():
			body += getBody(submessage)
	else:
		body = message.get_payload()
	return body.encode("string_escape")
			
			

class MsgData:
	"Holds the data we want from an email"
	
	def __init__(self, msgString):
			#parse in string
			self.msg = email.message_from_string(msgString)

			#BUG: We can't get the email body because the varchar size limit seems to be 10,000, and some emails are larger than that. 
			#it stores as a list okay, but you can't encode a list...
			#maybe just limit the size of emails to 10K chars and pass what you can to SQL, but that is not the entirety of the info. 
			s = self.msg.get_payload()
			self.body = getBody(self.msg)
			
			self.msgID = self.msg.get("Message-ID")
			self.date = self.msg.get("Date")
			self.subject = self.msg.get("Subject")
			self.recip = self.msg.get("To")
			self.sender = self.msg.get("From")
			self.references = self.msg.get("References")
			self.inReplyTo = self.msg.get("In-Reply-To")
			self.refList = []
			
			#check and parse msg ID
			if self.msgID != None:
				self.msgID = re.sub('.*<', '', self.msgID)
	        		self.msgID = re.sub('>.*', '', self.msgID)
			elif self.subject != None:
				raise DataError, 'No message id, subject is: ' + self.subject
			else: raise DataError, 'No message id'
			
			#check and parse date
			if self.date != None:
				self.parseDate()
			else: raise DataError, 'No date field!, MessageID: ' + self.msgID
			
				
			#parse subject
			
			if self.subject != None:
				#insert escape characters
				self.subject = self.subject.encode("string_escape")
			else: self.subject = ""
			
			#check sender and parse out sender email and name
			if self.sender != None:
				#lowercase everything
				self.sender = string.lower(self.sender)
				#change ' at ' to '@' - the gnucash mailing list tries to fool spammers this way
				self.sender = re.sub(' at ', '@', self.sender)
				#get at the string inside the <> (this is the email)
				self.senderEmail = re.sub('.*<', '', self.sender)
				self.senderEmail = re.sub('>.*', '', self.senderEmail)
				#take out name in parenthesis
				self.senderEmail = re.sub('\s*\(.*\)\s*', '', self.senderEmail)
				self.senderEmail = self.senderEmail.replace("'", '')
				if self.senderEmail.find('@') == -1:
					raise DataError, 'Invalid Email!, email: ' + self.senderEmail
				
				#take out email
				self.senderName = re.sub('[^@^\s]+@[^@^\s]+', '', self.sender)
				#take out special characters
				self.senderName = re.sub('\'|\"|\s+|<|>|\(|\)', ' ', self.senderName)
				#self.senderName = re.sub('[^[a-zA-Z\.-]]', ' ', self.senderName)
				#take out these odd URLs that sometimes go in, useless info.
				self.senderName = re.sub('=.*', '', self.senderName)
				#trim leading whitespace and trailing whitespace
				self.senderName = re.sub('^\s*|\s$', '', self.senderName)

			else: raise DataError, 'No sender line, message ID: ' + self.msgID
			
			#check recipient list and parse out emails and names
			if self.recip != None:
				self.recipEmail =  re.sub('.*<', '', self.recip)
				self.recipEmail =  re.sub('>.*', '', self.recipEmail)
				#change ' at ' to '@' - the gnucash mailing list tries to fool spammers this way
				self.recipEmail = re.sub(' at ', '@', self.recipEmail)
			#else: raise DataError, 'No recipient line, message ID:' + self.msgID
			else: self.recipEmail = self.recip = ''
			
			

			#check references
			if self.references != None:
				#trim out any prefix
				s = re.sub('.*<', '<', self.references)
				#kill any " or '
				s = re.sub("\'|\"", '', s)
				#replace <>'s with whitespace then trim edge white space
				s = re.sub('<|>',' ', s)
				s = re.sub('(^\s*)|(\s*\Z)','', s)
				#make a list of references, then reverse so most relevent is first
				self.refList = re.split('\s+', s)
				self.refList.reverse()
			#if we don't have references, use the reply to field
			#sometimes the replyto in a header isn't a msg id, so we don't use it if we don't have to
			elif self.inReplyTo != None:
				#strip off <>'s
				ref = re.sub('(.*<)|(>.*)', '', self.inReplyTo)
				#kill any " or '
				ref = re.sub("\'|\"", '', ref)
				refmo = re.search('[^<> \t\n\r\f\v]*@[^<> \t\n\r\f\v]*', ref)
				if refmo:
					#self.inReplyTo = refmo.group(0)
					self.refList.append(refmo.group(0))
				
				


	def hasReference(self):
		if len(self.refList) > 0:
			return True
		else: 
			return False
	
	def printData(self):
		print "Msg ID =", self.msgID
		print "Date =", self.date
		print "Subject =", self.subject
		print "To =", self.recip
		#print "ToEmail =", self.recipEmail
		print "From =", self.sender
		#print "FromEmail =", self.senderEmail
		if self.hasReference():
			print "first ref: ", self.refList[0]
		print "print first ref ="
	
	def parseDate(self):		
		#assumes self.date is the line from the email header after 'Date:'
		
		#check for weird characters in ISO
		mo = re.search('\?ISO', self.date)
		if mo:
			raise DataError, 'date had invalid characters'
		#first lets try the normal method, this works on apache, gnome, gtk, at least
		self.dateTuple = email.Utils.parsedate(self.date)
		#see if dateTuple is parseable
		try:
			self.calcSQLTimestamp()
		except DataError, de:
			#obviously we need to do some more work
			pass
		else:
			#we're good, lets get on with the show
			return
		
		#okay so lets try a hack for gnucash's (eventual) date format
		datetime = re.split('\s+', self.date)
		newdate = datetime[0] +" "+ datetime[2] +" "+ datetime[1] +" "+ datetime[4] +" "+ datetime[3]
		self.dateTuple = email.Utils.parsedate(newdate)
		#see if dateTuple is parseable
		try:
			self.calcSQLTimestamp()
		except DataError, de:
			#obviously we need to do some more work
			raise DataError, de.value
		else:
			#we're good, lets get on with the show
			return
	def calcTimeZone(self):
			#assume first field
			#tz = re.sub(".*:[0-9][0-9] ",'', self.date)
			#tz = re.sub(" .*",'', tz)
			#convert to sql readable timezone via a hashtable
			#tz = timeZoneTable.getTZ(tz)
		#assume self.date is filled in, and timezone is the last part of it
		tz = re.split("\s+", self.date)
		#take off newline
		tz[-1] = re.sub('\n', '', tz[-1])
		
		#test to see if the last part of the date is a validish timezone
		
		#test for +/- xxxx in last and second to last parts of the date line
		mo = re.match('\+[0-9]|-[0-9]', tz[-1])
		if mo:
			#convert to sql readable timezone via a hashtable
			self.timeZone = timeZoneTable.getTZ(tz[-1])
			return
		if(len(tz) > 1):
			mo = re.match('\+[0-9]|-[0-9]', tz[-2])
			if mo:
				#convert to sql readable timezone via a hashtable
				self.timeZone = timeZoneTable.getTZ(tz[-2])
				return
		#test for at least 2 uppercase letters, hopefully this isn't a day/month
		mo = re.search('[A-Z]{2,4}', tz[-1])
		if mo:
			#convert to sql readable timezone via a hashtable
			self.timeZone = timeZoneTable.getTZ(mo.group(0))
			return

		#doesn't seem to be a timezone given, even one we don't recognize
		else:
			self.timeZone = '+0000'
		
	def calcSQLTimestamp(self):
		"return date in a form sql can read"
		#assumes date and dateTuple are filled in
		self.calcTimeZone()
		#if timeZone is null then make it empty string
		if self.timeZone == None:
			self.timeZone = ""
			
		try:
			self.sqlTimestamp = time.strftime("%Y-%m-%d %H:%M:%S",self.dateTuple) + self.timeZone
		except ValueError:
			#date format is whack
			raise DataError, 'Date is in an unreadable format. \nformat: ' + self.date + " \ntuple: "+str(self.dateTuple)
		except TypeError:
			#not sure what happened
			raise DataError, 'Date could not be parsed. \nformat: ' + self.date + " \ntuple: "+str(self.dateTuple)
			
	def isValid(self):
		"Checks integrity of data, (missing data mostly)"
		if (self.msgID == None or self.subject == None or self.recip == None or self.sender == None or self.date == None):
			return False
		return True

def FileToMsgList(inputFile):
	"split a mbox file into a list of strings containing one message each"
	msgList = []
	while True:
		
		s = inputFile.readline()
		if s == "":
			return msgList
		elif re.match("From ", s):
			#we found the first message
			msg = s
			break
	#now we have the 'From' line of an email read into msg
	while True:
	
		s = inputFile.readline()
		if   s == "":
			#end of file
			msgList.append(msg)
			return msgList
		elif s.startswith("From "):
			#we found the next msg's 'From' line
			msgList.append(msg)
			msg = s
		else:
			#we are in the middle of the same msg
			msg = msg + s

def createTables(db):
	"drop any existing mail tables and create new ones"
	try:
		print "cleaning Email tables..."
		dropTablesQuery = getSqlFileText("dropEmailTables.sql")
		db.query(dropTablesQuery)
		print "...OK!"
	except _pg.ProgrammingError:
			print "could not drop tables"
	try:
		print "creating Email tables..."
		createTablesQuery = getSqlFileText("createEmailTables.sql")
		db.query(createTablesQuery)
		print "...OK!"
	except _pg.ProgrammingError:
			print "could not create tables"

#convert the string to ascii format (could be unicode) and replace non-ascii
#characters to a ?
def toAscii(s):
	return unicode(s, errors="ignore")
		
def insertIntoSQL(db, msgDataList, listName):
	"parse a list of msg data classes into a sql table"
			
	#loop msgData classes and parse out all emails and put into alias table
	for msgData in msgDataList:	
		select = db.query("SELECT AliasID FROM Aliases where email = '" + toAscii(msgData.senderEmail) +"'")
		if select.getresult() == []: #we didn't find the entry
			#insert new entry
			s = "INSERT INTO Aliases (email, name) VALUES('"+ toAscii(msgData.senderEmail) + "','" + toAscii(msgData.senderName) +"')"
			#print s
			ecquery(db,s, 'sender name')
	
	#loop msgData classes and parse out messages
	for msgData in msgDataList:
		#constraint of subject length in sql
		if len(msgData.subject) >= 512:
			msgData.subject = msgData.subject[0:511]
		s = "INSERT INTO Messages (MessageID, SenderAliasID, DateTime, Body, Subject, List) VALUES('" + toAscii(msgData.msgID)
		s = s + "', (select AliasID from Aliases where Email = '"+ toAscii(msgData.senderEmail) +"'),'" + msgData.sqlTimestamp +"', '" + toAscii(msgData.body) + "','" + toAscii(msgData.subject) + "','"+ listName +"' )"
		#print s
		#ecquery(db,s, 'msg id or body')
		ecquery(db,s, msgData.msgID)
		#error handling because pygresql has only one kind of exception type
	

def insertReferencesIntoSQL(db, referencesDict):
	"""update the table with references.  Since all messages have already been
	put into the database, if a reference isn't in the table then it doesn't
	still exist in the archive."""
	referenceUpdates = 0
	for msgID, refIDs in referencesDict.items():
		#only store the first reference that exists in the database
		for refID in refIDs:
			if referencesDict.has_key(refID):
				#store the reference id and the recipient alias for this message
				query = "UPDATE Messages SET ReferenceID = '" + refID + "', RecipAliasId = (select SenderAliasID from Messages where MessageID = '" + refID + "')"
				query += "where MessageId = '" + msgID + "'"
				ecquery(db, query, 'reply id')
				referenceUpdates += 1
				if referenceUpdates % 1000 == 0:
					print "Updated references for %i messages" % referenceUpdates
				break
	print "Updated references for %i messages" % referenceUpdates
	
		
#error check query
def ecquery(db, s, part):
	try:
		db.query(s)
	except _pg.ProgrammingError, e:
		mo1 = re.search("violates unique", str(e))
		if mo1:
			#dup message, ignore this error
			print "DUPLICATE MESSAGE, ommiting message", str(e)
			return
		mo2 = re.search("UNICODE", str(e))
		if mo2:
			#unicode error
			print "UNICODE ERROR:", part, "ommiting message", str(e)
			return
		mo2 = re.search("null value in column \"senderaliasid\" violates not-null constraint", str(e))
		if mo2:
			print "NO ALIAS RECORDED, ommiting message", str(e)
			return
		mo2 = re.search("insert or update on table \"messages\" violates foreign key constraint \"messages_referenceid_fkey\"", str(e))
		if mo2:
			print "RE MSGID NONEXISTANT, ommiting message"
			print s
			return
		print "error inserting:",s
		
		print e
		sys.exit(2)		

def parseMbox(db, mboxFolder, listName):
	#print "Processing File", arg, "..."
	
	#Get Files in mboxFolder
	s = 'ls ' + mboxFolder
	mboxFiles = popen2.popen2(s)[0] #mbox is a file handle to a list of files in dir
	mboxFileList =  mboxFiles.read().split() #mboxFileList is a list of filenames

	msgList = []
	totalCount = 0
	totalParsed = 0
	msgFailures = 0
	# a dictionary where key is messageid for a message and the value
	# is the list of messageid's that this e-mail references
	referencesDict = {}

	
	#To test the script on just a few messages:
	#mboxFile = mboxFileList[1]

	# go through the process of splitting messages, parsing messages, and committing them
	# to the database on a file by file basis (month by month)
	for mboxFile in mboxFileList:
		msgList = []
		if mboxFile == 'CVS':
			continue	
		fileName = os.path.join(mboxFolder, mboxFile)
		if not os.path.isfile(fileName):
			continue
		inputFile = open(fileName)
		#convert to a list of strings
		ml = FileToMsgList(inputFile)
		msgList.extend(ml)
		print fileName, "contains", len(ml), "messages..."
		totalCount += len(ml)
		inputFile.close()
		
		#store the messages from the input file in a list of msgdata classes
		msgDataList = []
		for msg in msgList:
			try:
				msgData = MsgData(msg)
				msgDataList.append(msgData)
				#store the list of references into the references dictionary
				referencesDict[msgData.msgID] = msgData.refList
				#show how much progress is being made
				totalParsed += 1
				if totalParsed % 1000 == 0:
					print "*** Parsed %i messages ***" % totalParsed
			except email.Errors.MessageError, e:
				print "     ***error on parsing message***"
				print "    ",e
				msgFailures += 1
			except DataError, e:
				print "     ***error on parsing message***"
				print "    ",e.value
				msgFailures += 1
			#don't let anything stop the script, just move on to the next message
			except Exception, e:
				print "     ***error on parsing message***"
				print "    ",e
				msgFailures += 1
				
		#put message data into SQL
		insertIntoSQL(db, msgDataList, listName)
		
	# Now that every message has been put into the database, we can update the references
	# and not violate the foreign key constraint
	insertReferencesIntoSQL(db, referencesDict)

	print "There are a total of %i messages" % totalCount
	print "failed to parse %i messages" % msgFailures
	
	print "Parsed", totalParsed, "messages successfully"
	s = "select count(*) from messages"
	pgo = db.query(s)
	successful = pgo.getresult()[0][0]
	print "Put", successful, "messages into database"
	print 100.0*successful/(totalCount + .0001), "percent of the messages submited error free into database"
